from core.function import Function
from core.cuda import cuda
from utils.common import to_variable
import utils

class SumTo(Function):
    def __init__(self, shape):
        """
        初始化 SumTo 函数。

        Parameters:
            shape (tuple): 目标形状。
        """
        self.shape = shape

    def forward(self, x):
        """
        执行 SumTo 操作的前向传播。

        Parameters:
            x (Variable): 输入变量。

        Returns:
            Variable: 操作结果。
        """
        x = cuda.to_array(x)
        self.x_shape = x.shape

        y = self._sum_to(x, self.shape)
        return y

    def backward(self, gy):
        """
        执行 SumTo 操作的反向传播。

        Parameters:
            gy (Variable): 相对于前向传播结果的梯度。

        Returns:
            Variable: 相对于输入变量的梯度。
        """
        gx = utils.functions_collect.broadcast_to(gy, self.x_shape)
        return gx

    def _sum_to(self, x, shape):
        """
        沿着轴求和以产生给定形状的数组。

        Args:
            x (ndarray): 输入数组。
            shape (tuple): 目标形状。

        Returns:
            ndarray: 输出数组，具有指定的形状。
        """
        ndim = len(shape)
        lead = x.ndim - ndim
        lead_axis = tuple(range(lead))

        axis = tuple([i + lead for i, sx in enumerate(shape) if sx == 1])
        y = x.sum(lead_axis + axis, keepdims=True)
        if lead > 0:
            y = y.squeeze(lead_axis)
        return y

    def sum_to(self, x):
        """
        执行 SumTo 操作。

        Parameters:
            x (Variable): 输入变量。

        Returns:
            Variable: 操作结果。
        """
        if x.shape == self.shape:
            return to_variable(x)
        return self(x)


